from tests.pytests.helpers.codegen_helpers import compile_chunk, assert_code_contains
from typing import List, Optional


def test_optional_and_list_mapping():
    def test_func(xs: List[int], y: Optional[int]) -> int:
        if y is not None:
            return xs[0] + y
        return xs[0]

    block = compile_chunk(test_func)
    assert_code_contains(
        block,
        """
        int test_func(const mc::array<int>& xs, const std::optional<int>& y) {
            if (y.has_value()) {
                return mcpy::list_getitem(xs, 0) + y;
            }
            return mcpy::list_getitem(xs, 0);
        }
        """,
    )


def test_typevar_generic_function():
    """测试使用 TypeVar 的泛型函数"""
    code = """
from typing import TypeVar

T = TypeVar('T')

def identity(x: T) -> T:
    return x
"""
    cpp = compile_chunk(code)
    expected = """
    template <typename T>
    T identity(const T& x) {
        return x;
    }
    """
    assert_code_contains(cpp, expected)


def test_typevar_generic_function_with_multiple_params():
    """测试使用 TypeVar 的泛型函数（多个参数）"""
    code = """
from typing import TypeVar

T = TypeVar('T')

def first(x: T, y: T) -> T:
    return x
"""
    cpp = compile_chunk(code)
    expected = """
    template <typename T>
    T first(const T& x, const T& y) {
        return x;
    }
    """
    assert_code_contains(cpp, expected)


def test_typevar_generic_function_multiple_typevars():
    """测试使用多个 TypeVar 的泛型函数"""
    code = """
from typing import TypeVar

T = TypeVar('T')
U = TypeVar('U')

def pair_first(x: T, y: U) -> T:
    return x
"""
    cpp = compile_chunk(code)
    expected = """
    template <typename T, typename U>
    T pair_first(const T& x, const U& y) {
        return x;
    }
    """
    assert_code_contains(cpp, expected)


def test_non_generic_function_with_typevar_defined():
    """测试定义了 TypeVar 但未使用的函数"""
    code = """
from typing import TypeVar

T = TypeVar('T')

def add(x: int, y: int) -> int:
    return x + y
"""
    cpp = compile_chunk(code)
    # 应该不生成模板
    expected = """
    int add(int x, int y) {
        return x + y;
    }
    """
    assert_code_contains(cpp, expected)


def test_typevar_container_mutation():
    """测试模板函数中容器参数的修改"""
    code = """
from typing import TypeVar

T = TypeVar('T')

def append_to_list(container: list[T], value: T) -> None:
    container.append(value)
"""
    cpp = compile_chunk(code)
    expected = """
    template <typename T>
    void append_to_list(mc::array<T>& container, const T& value) {
        container.push_back(value);
    }
    """
    assert_code_contains(cpp, expected)


def test_typevar_parameter_mutation():
    """测试模板参数被修改的情况"""
    code = """
from typing import TypeVar

T = TypeVar('T')
U = TypeVar('U')

def add_to(accumulator: T, value: U) -> None:
    accumulator += value
"""
    cpp = compile_chunk(code)
    expected = """
    template <typename T, typename U>
    void add_to(T& accumulator, const U& value) {
        accumulator += value;
    }
    """
    assert_code_contains(cpp, expected)
